refactor(flowcontrol): single-ended SafeQueue backed by container/heap#1690
Conversation
The dispatch SafeQueue was double-ended, but nothing consumed PeekTail other than the flowQueueAccessor passthrough itself: dispatch needs only the head, and eviction uses an independent ordering. With PeekTail gone the queue is single-ended, so the bespoke max-min heap that gave O(1) access to both ends is no longer justified. - Remove the unused PeekTail from the SafeQueue contract and all implementers, mocks, and tests; rename PeekHead -> Peek now that there is a single end. - Replace the hand-rolled max-min heap (maxminheap.go) with a standard library container/heap priority queue (priorityqueue.go), eliminating the custom bubble-up/down logic and the latent bug class it carried. - Drop the side handle map: membership is validated by index identity (a *heapItem lives in exactly one queue), preserving the contract's Remove error semantics with less state. O(log n) targeted removal is retained via per-item index tracking. - Rename the registered queue MaxMinHeap -> PriorityQueue. Signed-off-by: Luke Van Drie <lukevandrie@google.com>
github-actions
Bot
added
size/XXL
| pq.mu.Lock() | ||
| defer pq.mu.Unlock() |
There was a problem hiding this comment.
The statements that need locking are 189 through 193, perhaps the Lock and Unlock calls should wrap those statements
There was a problem hiding this comment.
Partially applied. I hoisted the nil check and the type assertion out (they only inspect the caller's handle). But the IsInvalidated() check has to stay inside the lock: isInvalidated (and index) are written under this same mutex by Remove/Cleanup/Drain, so reading them lock-free would be a data race.
So the locked region is now: IsInvalidated → identity check → heap.Remove → invalidate.
| hi.index = -1 | ||
| pq.byteSize.Add(^hi.item.OriginalRequest().ByteSize() + 1) // Atomic subtraction. | ||
| } else { | ||
| itemsToKeep = append(itemsToKeep, hi) |
There was a problem hiding this comment.
As itemsToKeep is pre-allocated to be the size of the current heap, which will be more than sufficient even if no items are removed by the predicate. Why not have a counter and do simple assignments instead of calling append here?
There was a problem hiding this comment.
Nice call; I switched to in-place compaction which drops the second slice allocation entirely. Added a small tail-nil loop so the removed *heapItems aren't retained by the vacated capacity.
| pq.mu.Lock() | ||
| defer pq.mu.Unlock() | ||
|
|
||
| drainedItems := make([]flowcontrol.QueueItemAccessor, len(pq.heap.items)) | ||
| for i, hi := range pq.heap.items { | ||
| drainedItems[i] = hi.item | ||
| hi.Invalidate() | ||
| hi.index = -1 | ||
| } | ||
|
|
||
| pq.heap.items = make([]*heapItem, 0) | ||
| pq.byteSize.Store(0) |
There was a problem hiding this comment.
| pq.mu.Lock() | |
| defer pq.mu.Unlock() | |
| drainedItems := make([]flowcontrol.QueueItemAccessor, len(pq.heap.items)) | |
| for i, hi := range pq.heap.items { | |
| drainedItems[i] = hi.item | |
| hi.Invalidate() | |
| hi.index = -1 | |
| } | |
| pq.heap.items = make([]*heapItem, 0) | |
| pq.byteSize.Store(0) | |
| pq.mu.Lock() | |
| old := pq.heap.items | |
| pq.heap.items = make([]*heapItem, 0) | |
| pq.byteSize.Store(0) | |
| pq.mu.Unlock() | |
| drainedItems := make([]flowcontrol.QueueItemAccessor, len(old)) | |
| for i, hi := range old { | |
| drainedItems[i] = hi.item | |
| hi.Invalidate() | |
| hi.index = -1 | |
| } |
There was a problem hiding this comment.
I'd like to push back on this one. Moving the loop outside the lock makes hi.Invalidate() / hi.index = -1 race with a concurrent Remove(handle). Even after the slice is swapped out, a caller can still hold a handle to one of the drained items and call Remove, which reads hi.isInvalidated and hi.index under the lock. Drain mutating those same fields lock-free is a data race.
If we ever want to shorten this O(n) section, the clean way would be making isInvalidated/index atomic so invalidation can go lock-free, but that's a broader change I'd rather not fold into this PR.
There was a problem hiding this comment.
Also, Drain is only ran on the graceful shutdown path. I don't think we need to be concerned with optimizing this method.
Addresses review comments from @shmuelk on llm-d#1690: - Correct the new-file license headers to "2026 The llm-d Authors". - Clarify the PriorityQueue doc: ordering is maintained by an internal container/heap (the queue does not itself implement heap.Interface). - Narrow the Add and Remove critical sections to the shared heap mutation. Precondition checks (nil, type assertion) and the atomic byteSize update move outside the lock; the isInvalidated and index reads stay inside it, since those fields are written under the lock by Remove/Cleanup/Drain and reading them lock-free would race. - Compact Cleanup survivors in place instead of allocating a second slice, clearing the vacated tail so removed items are not retained. Signed-off-by: Luke Van Drie <lukevandrie@google.com>
3 participants
What type of PR is this?
/kind cleanup
What this PR does / why we need it:
The flow-control
SafeQueuewas double-ended (PeekHead+PeekTail), backed by a hand-rolled max-min heap. @RishabhSaini identified a latent bug in that bespoke data structure (#1689). In practice nothing consumesPeekTail, the only reference was theflowQueueAccessorpassthrough itself. Dispatch (EDF/SLO ordering) needs only the head,and eviction uses an independent ordering, so a double-ended dispatch queue buys nothing.
This makes the queue single-ended and swaps the custom heap for the standard library:
PeekTailfromQueueInspectionMethodsand every implementer, mock, and test.PeekHead->Peeknow that there is a single end.maxminheap.gowithpriorityqueue.go, acontainer/heap-backed priority queue (PriorityQueue). This deletes the custom sift logic and the latent bug class with it, while keeping O(log n) targeted removal via per-item index tracking.Removevalidates membership by index identity, preserving the existing error contract with less state.MaxMinHeap->PriorityQueue.container/heapwith index tracking is already the established pattern inflowcontrol/eviction/queue.go, so this aligns the dispatch queue with existing style. FCFS /ListQueuebehavior is unchanged.Which issue(s) this PR fixes:
N/A, but a more aggressive refactoring that resolves the issue surfaced in #1689.
Release note (write
NONEif no user-facing change):